from typing import Callable

import numpy as np
from scipy.spatial.transform import Rotation as R

from dora import DoraStatus
from transform_utils import state2SE3, sE32states

# Kalman Filter
KALMAN_FILTER_NUM_GNSS_INIT_SAMPLES = 10

# Vehicle
VEHICLE_SENSOR_POSITION = np.array(
    [0, 0, 1.5, 0, 0, 0]
)  # assume imu and gnss are mounted at same location
VEHICLE_DT = 1 / 20.0  # s

# Debug
DEBUG = False

if DEBUG:
    import matplotlib.pyplot as plt
    from collections import deque


def angle_normalize(a):
    """Normalize angles to lie in range -pi < a[i] <= pi."""
    a = np.remainder(a, 2 * np.pi)
    a = np.where(a <= -np.pi, a + 2 * np.pi, a)
    a = np.where(a > np.pi, a - 2 * np.pi, a)
    return a


def skew_symmetric(v):
    """Skew symmetric form of a 3x1 vector."""
    return np.array(
        [[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=np.float64
    )


def omega(w, delta_t):
    """
    Purturbation locally at the SO3 manifold
    (right quaternion product matrix)
    """
    theta = w * delta_t
    q = Quaternion(axis_angle=angle_normalize(theta)).to_numpy()
    qw = q[0]
    qv = q[1:]

    om1 = np.zeros((4, 4))
    om1[0, 1:] = -qv
    om1[1:, 0] = qv
    om1[1:, 1:] = -skew_symmetric(qv)
    om = np.identity(4) * qw + om1

    return om


class Quaternion:
    def __init__(self, w=1.0, x=0.0, y=0.0, z=0.0, axis_angle=None, euler=None):
        """
        Allow initialization with explicit quaterion wxyz, axis-angle, or Euler XYZ (RPY) angles.

        :param w: w (real) of quaternion.
        :param x: x (i) of quaternion.
        :param y: y (j) of quaternion.
        :param z: z (k) of quaternion.
        :param axis_angle: Set of three values from axis-angle representation, as list or [3,] or [3,1] np.ndarray.
                           See C2M5L2 for details.
        :param euler: Set of three XYZ Euler angles.
        """
        if axis_angle is None and euler is None:
            self.w = w
            self.x = x
            self.y = y
            self.z = z
        elif euler is not None and axis_angle is not None:
            raise AttributeError("Only one of axis_angle or euler can be specified.")
        elif axis_angle is not None:
            if (
                not (type(axis_angle) == list or type(axis_angle) == np.ndarray)
                or len(axis_angle) != 3
            ):
                raise ValueError("axis_angle must be list or np.ndarray with length 3.")
            axis_angle = np.array(axis_angle)
            norm = np.linalg.norm(axis_angle)
            self.w = np.cos(norm / 2)
            if norm < 1e-50:  # to avoid instabilities and nans
                self.x = 0
                self.y = 0
                self.z = 0
            else:
                imag = axis_angle / norm * np.sin(norm / 2)
                self.x = imag[0].item()
                self.y = imag[1].item()
                self.z = imag[2].item()
        else:
            roll = euler[0]
            pitch = euler[1]
            yaw = euler[2]

            cy = np.cos(yaw * 0.5)
            sy = np.sin(yaw * 0.5)
            cr = np.cos(roll * 0.5)
            sr = np.sin(roll * 0.5)
            cp = np.cos(pitch * 0.5)
            sp = np.sin(pitch * 0.5)

            # Fixed frame
            self.w = cr * cp * cy + sr * sp * sy
            self.x = sr * cp * cy - cr * sp * sy
            self.y = cr * sp * cy + sr * cp * sy
            self.z = cr * cp * sy - sr * sp * cy

            # Rotating frame
            # self.w = cr * cp * cy - sr * sp * sy
            # self.x = cr * sp * sy + sr * cp * cy
            # self.y = cr * sp * cy - sr * cp * sy
            # self.z = cr * cp * sy + sr * sp * cy

    def __repr__(self):
        return "Quaternion (wxyz): [%2.5f, %2.5f, %2.5f, %2.5f]" % (
            self.w,
            self.x,
            self.y,
            self.z,
        )

    def to_axis_angle(self):
        t = 2 * np.arccos(self.w)
        return np.array(t * np.array([self.x, self.y, self.z]) / np.sin(t / 2))

    def to_mat(self):
        v = np.array([self.x, self.y, self.z]).reshape(3, 1)
        return (
            (self.w**2 - np.dot(v.T, v)) * np.eye(3)
            + 2 * np.dot(v, v.T)
            + 2 * self.w * skew_symmetric(v)
        )

    def to_euler(self):
        """Return as xyz (roll pitch yaw) Euler angles."""
        roll = np.arctan2(
            2 * (self.w * self.x + self.y * self.z), 1 - 2 * (self.x**2 + self.y**2)
        )
        pitch = np.arcsin(2 * (self.w * self.y - self.z * self.x))
        yaw = np.arctan2(
            2 * (self.w * self.z + self.x * self.y), 1 - 2 * (self.y**2 + self.z**2)
        )
        return np.array([roll, pitch, yaw])

    def to_numpy(self):
        """Return numpy wxyz representation."""
        return np.array([self.w, self.x, self.y, self.z])

    def normalize(self):
        """Return a (unit) normalized version of this quaternion."""
        norm = np.linalg.norm([self.w, self.x, self.y, self.z])
        return Quaternion(self.w / norm, self.x / norm, self.y / norm, self.z / norm)

    def quat_mult_right(self, q, out="np"):
        """
        Quaternion multiplication operation - in this case, perform multiplication
        on the right, that is, q*self.

        :param q: Either a Quaternion or 4x1 ndarray.
        :param out: Output type, either np or Quaternion.
        :return: Returns quaternion of desired type.
        """
        v = np.array([self.x, self.y, self.z]).reshape(3, 1)
        sum_term = np.zeros([4, 4])
        sum_term[0, 1:] = -v[:, 0]
        sum_term[1:, 0] = v[:, 0]
        sum_term[1:, 1:] = -skew_symmetric(v)
        sigma = self.w * np.eye(4) + sum_term

        if type(q).__name__ == "Quaternion":
            quat_np = np.dot(sigma, q.to_numpy())
        else:
            quat_np = np.dot(sigma, q)

        if out == "np":
            return quat_np
        elif out == "Quaternion":
            quat_obj = Quaternion(quat_np[0], quat_np[1], quat_np[2], quat_np[3])
            return quat_obj

    def quat_mult_left(self, q, out="np"):
        """
        Quaternion multiplication operation - in this case, perform multiplication
        on the left, that is, self*q.

        :param q: Either a Quaternion or 4x1 ndarray.
        :param out: Output type, either np or Quaternion.
        :return: Returns quaternion of desired type.
        """
        v = np.array([self.x, self.y, self.z]).reshape(3, 1)
        sum_term = np.zeros([4, 4])
        sum_term[0, 1:] = -v[:, 0]
        sum_term[1:, 0] = v[:, 0]
        sum_term[1:, 1:] = skew_symmetric(v)
        sigma = self.w * np.eye(4) + sum_term

        if type(q).__name__ == "Quaternion":
            quat_np = np.dot(sigma, q.to_numpy())
        else:
            quat_np = np.dot(sigma, q)

        if out == "np":
            return quat_np
        elif out == "Quaternion":
            quat_obj = Quaternion(quat_np[0], quat_np[1], quat_np[2], quat_np[3])
            return quat_obj


class ExtendedKalmanFilter:
    """
    Extended Kalman Filter
    """

    def __init__(self):
        # State
        self.pos = np.zeros((3, 1))  # position: [x, y, z]
        self.vel = np.zeros((3, 1))  # velocity: [vx, vy, vz]
        self.quat = np.zeros((4, 1))  # orientation: [qw, qx, qy, qz]

        # State covariance (we use global angluar errors for orientation)
        self.p_cov = np.zeros((9, 9))

        # Gravity
        self.g = np.array([0, 0, -9.81]).reshape(-1, 1)

        # Time step
        self.dt = VEHICLE_DT

        # Sensor measurement variances
        self.var_imu_acc = np.diag(
            [
                1.0e-1,  # accel_x
                1.0e-1,  # accel_y
                1.0e-1,  # accel_z
            ]
        )
        self.var_imu_gyro = np.diag(
            [
                1.0e-1,  # gyro_x
                1.0e-1,  # gyro_y
                1.0e-1,  # gyro_z
            ]
        )
        self.var_measurement = np.diag(
            [
                5.0e-1,  # gnss_x
                5.0e-1,  # gnss_y
                1.0e-5,  # gnss_z
                1.0e-5,  # yaw
                1.0e-1,  # speed
            ]
        )

        # Motion model noise Jacobian
        self.l_jac = np.zeros((9, 6))
        self.l_jac[3:, :] = np.eye(6)

        # Measurement model Jacobian
        self.h_jac = np.zeros([5, 9])  # gnss_x, gnss_y, gnss_z
        self.h_jac[:3, :3] = np.eye(3)  # yaw
        self.h_jac[3, 8] = 1  # speed

        # Status
        self.has_initialized = False

    def set_initial_state(self, init_state):
        """
        init_state = [x, y, z, vx, vy, vz, roll, pitch, yaw]
        """
        if self.has_initialized:
            return False

        else:
            self.pos[:, 0] = np.array(init_state[:3])
            self.vel[:, 0] = np.array(init_state[3:6])
            self.quat[:, 0] = Quaternion(euler=init_state[6:9]).to_numpy()

            init_estimate_covariance = np.diag(
                [
                    # x, y, z,
                    5.0e-1,
                    5.0e-1,
                    1.0e-5,
                    # vx, vy, vz,
                    1.0e-3,
                    1.0e-3,
                    1.0e-3,
                    # angle_x, angle_y, angle_z
                    1.0e-1,
                    1.0e-1,
                    1.0e-5,
                ]
            )
            self.p_cov = init_estimate_covariance
            self.has_initialized = True

            return True

    def get_position(self):
        return self.pos.reshape(-1)

    def get_velocity(self):
        return self.vel.reshape(-1)

    def get_speed(self):
        return np.linalg.norm(self.vel)

    def get_orientation(self, to_euler=False):
        if to_euler:
            qw, qx, qy, qz = self.quat.flatten()
            euler_angles = R.from_quat([qx, qy, qz, qw]).as_euler("xyz", degrees=False)
            return np.array(euler_angles)
        else:
            return self.quat.reshape(-1)

    def predict_state(self, imu):
        """
        imu = [accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z]
        """
        # IMU accelerometer and gyroscope readings
        imu_accel = np.array(imu[:3], dtype=np.float64).reshape(-1, 1)
        imu_gyro = np.array(imu[3:6], dtype=np.float64).reshape(-1, 1)

        # Update state with IMU data
        R = Quaternion(*self.quat).to_mat()
        self.pos = (
            self.pos
            + self.dt * self.vel
            + 0.5 * (self.dt**2) * (R @ imu_accel + self.g)
        )
        self.vel = self.vel + self.dt * (R @ imu_accel + self.g)
        self.quat = omega(imu_gyro, self.dt) @ self.quat

        # Linearize motion model and compute Jacobians
        F = np.eye(9)
        F[:3, 3:6] = np.eye(3) * self.dt
        F[3:6, 6:] = -skew_symmetric(R @ imu_accel) * self.dt

        # Propagate uncertainty
        Q = np.eye(6)
        Q[:3, :3] *= self.var_imu_acc * (self.dt**2)
        Q[3:, 3:] *= self.var_imu_gyro * (self.dt**2)
        self.p_cov = F @ self.p_cov @ F.T + self.l_jac @ Q @ self.l_jac.T

    def measurement_update(self, gnss, speed, heading):
        """
        gnss = [x, y, z], heading in radians
        """
        # GNSS readings in world frame
        gnss_meas = np.array(gnss, dtype=np.float64).reshape(-1, 1)

        # Compute quaternion innovation
        _, _, yaw_pred = self.get_orientation(to_euler=True)
        delta_yaw = angle_normalize(np.array([heading - yaw_pred]))[0]

        # Predicted speed from state
        speed_pred = np.linalg.norm(self.vel)

        # Update measurement model Jacobian
        if speed_pred > 0:
            self.h_jac[4, 3:6] = self.vel.flatten() / speed_pred
        else:
            self.h_jac[4, 3:6] = np.zeros(3)

        # Augment the measurement
        z_meas = np.vstack((gnss_meas, delta_yaw, speed))

        # Compute Kalman gain
        K = (
            self.p_cov
            @ self.h_jac.T
            @ (
                np.linalg.inv(
                    self.h_jac @ self.p_cov @ self.h_jac.T + self.var_measurement
                )
            )
        )

        # Compute error state
        delta_x = K @ (z_meas - np.vstack((self.pos, 0, speed_pred)))

        # Correct predicted state
        self.pos += delta_x[:3]
        self.vel += delta_x[3:6]
        delta_quat = Quaternion(axis_angle=angle_normalize(delta_x[6:9]))
        self.quat = delta_quat.quat_mult_left(self.quat)

        # Normalize quaternion
        self.quat = Quaternion(*self.quat).normalize().to_numpy().reshape(-1, 1)

        # Compute corrected covariance
        # self.p_cov = (np.eye(9) - K @ self.h_jac) @ self.p_cov
        self.p_cov = (np.eye(9) - K @ self.h_jac) @ self.p_cov @ (
            np.eye(9) - K @ self.h_jac
        ).T + K @ self.var_measurement @ K.T


class Operator:
    """
    Compute vehicle state using Extended Kalman Filter
    """

    def __init__(self):
        self.imu = np.array([])
        self.last_imu = np.array([])
        self.imu_init_xyz = []
        self.gnss = np.array([])
        self.gnss_init_xyz = []
        self.yaw = 0.0
        self.speed = 0.0
        # Extended Kalman Filter
        self.ekf = ExtendedKalmanFilter()

        if DEBUG:
            max_len = 50
            self.ekf_x_y_yaw_vel_history = deque(maxlen=max_len)
            self.sensor_x_y_yaw_vel_history = deque(maxlen=max_len)
            self.actual_x_y_yaw_vel_history = deque(maxlen=max_len)
            self.actual_position = np.array([])

            self.fig, self.axis = plt.subplots(4, 1, sharex=True)
            self.ekf_history_handle = []
            self.sensor_history_handle = []
            self.actual_history_handle = []
            for i in range(4):
                (ekf_handle,) = self.axis[i].plot(
                    [], [], "b-", linewidth=1.0, label="ekf" if i == 0 else None
                )
                (state_handle,) = self.axis[i].plot(
                    [], [], "r-", linewidth=1.0, label="sensor" if i == 0 else None
                )
                (actual_handle,) = self.axis[i].plot(
                    [], [], "g--", linewidth=1.0, label="actual" if i == 0 else None
                )
                self.ekf_history_handle.append(ekf_handle)
                self.sensor_history_handle.append(state_handle)
                self.actual_history_handle.append(actual_handle)

            self.axis[0].set_ylabel("x [m]")
            self.axis[1].set_ylabel("y [m]")
            self.axis[2].set_ylabel("yaw [rad]")
            self.axis[3].set_ylabel("speed [m/s]")
            self.axis[0].legend(fontsize=10)
            self.fig.tight_layout()

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):
        """Handle input.
        Args:
            dora_input["id"](str): Id of the input declared in the yaml configuration
            dora_input["data"] (bytes): Bytes message of the input
            send_output (Callable[[str, bytes]]): Function enabling sending output back to dora.
        """

        if dora_input["id"] == "imu":
            imu = np.frombuffer(dora_input["data"], np.float32)
            self.imu = imu[:-1]
            self.yaw = imu[-1] - np.pi / 2

            if not self.ekf.has_initialized:
                self.imu_init_xyz.append(imu[:3])

            if self.ekf.has_initialized and len(self.last_imu) != 0:
                self.ekf.predict_state(self.last_imu)
            self.last_imu = self.imu.copy()
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "speed":
            speed = np.frombuffer(dora_input["data"], np.float32)
            self.speed = speed[0]
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "position":
            if DEBUG:
                self.actual_position = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "gnss":
            gnss = np.frombuffer(dora_input["data"], np.float32)
            self.gnss = gnss

            if not self.ekf.has_initialized:
                self.gnss_init_xyz.append(gnss)

                if len(self.gnss_init_xyz) == KALMAN_FILTER_NUM_GNSS_INIT_SAMPLES or (
                    len(self.gnss_init_xyz) > 1 and abs(self.speed) > 0.1
                ):
                    gnss_init_xyz_avg = np.array(self.gnss_init_xyz).mean(axis=0)
                    # if len(self.imu_init_xyz) > 0:
                    #     imu_init_xyz_avg = np.array(self.imu_init_xyz).mean(axis=0)
                    #     init_roll = np.arctan2(imu_init_xyz_avg[1], imu_init_xyz_avg[2])
                    #     init_pitch = np.arctan2(
                    #         -imu_init_xyz_avg[0],
                    #         np.sqrt(
                    #             imu_init_xyz_avg[1] ** 2 + imu_init_xyz_avg[2] ** 2
                    #         ),
                    #     )
                    # else:
                    #     init_roll = 0.0
                    #     init_pitch = 0.0
                    init_roll = 0.0
                    init_pitch = 0.0

                    init_state = np.array(
                        # x, y, z, vx, vy, vz, roll, pitch, yaw
                        [
                            gnss_init_xyz_avg[0],
                            gnss_init_xyz_avg[1],
                            gnss_init_xyz_avg[2],
                            0.0,
                            0.0,
                            0.0,
                            init_roll,
                            init_pitch,
                            self.yaw,
                        ],
                        dtype=np.float64,
                    )

                    self.ekf.set_initial_state(init_state)
            else:
                self.ekf.measurement_update(
                    gnss=self.gnss, speed=self.speed, heading=self.yaw
                )

                x, y, z = self.ekf.get_position()
                roll, pitch, yaw = self.ekf.get_orientation(to_euler=True)

                # Convert to vehicle center
                vehicle_T_sensor = state2SE3(VEHICLE_SENSOR_POSITION)
                sensor_T_vehicle = np.linalg.inv(vehicle_T_sensor)
                # world_T_sensor = state2SE3([x, y, z, roll, pitch, yaw])
                world_T_sensor = state2SE3([x, y, self.gnss[2], roll, pitch, self.yaw])
                world_T_vehicle = world_T_sensor @ sensor_T_vehicle

                # Send output
                position_filtered = sE32states(
                    world_T_vehicle, to_quaternion=True
                ).astype(np.float32)

                send_output(
                    "position_filtered",
                    position_filtered.tobytes(),
                    {},
                )

                if DEBUG:
                    self.ekf_x_y_yaw_vel_history.append(
                        [x, y, angle_normalize(yaw), self.ekf.get_speed()]
                    )
                    self.sensor_x_y_yaw_vel_history.append(
                        [
                            self.gnss[0],
                            self.gnss[1],
                            angle_normalize(self.yaw),
                            self.speed,
                        ]
                    )
                    [
                        ego_x,
                        ego_y,
                        ego_z,
                        ego_rx,
                        ego_ry,
                        ego_rz,
                        ego_rw,
                    ] = self.actual_position
                    [ego_pitch, ego_roll, ego_yaw] = R.from_quat(
                        [ego_rx, ego_ry, ego_rz, ego_rw]
                    ).as_euler("xyz", degrees=False)
                    world_T_vehicle = state2SE3(
                        [ego_x, ego_y, ego_z, ego_roll, ego_pitch, ego_yaw]
                    )
                    world_T_sensor = world_T_vehicle @ vehicle_T_sensor
                    (actual_x, actual_y, _, _, _, actual_yaw) = sE32states(
                        world_T_sensor, to_quaternion=False
                    )
                    self.actual_x_y_yaw_vel_history.append(
                        [actual_x, actual_y, angle_normalize(actual_yaw), self.speed]
                    )
                    ekf_x_y_yaw_vel_history_np = np.array(self.ekf_x_y_yaw_vel_history)
                    sensor_x_y_yaw_vel_history_np = np.array(
                        self.sensor_x_y_yaw_vel_history
                    )
                    actual_x_y_yaw_vel_history_np = np.array(
                        self.actual_x_y_yaw_vel_history
                    )
                    for i in range(4):
                        self.ekf_history_handle[i].set_xdata(
                            range(len(ekf_x_y_yaw_vel_history_np))
                        )
                        self.ekf_history_handle[i].set_ydata(
                            ekf_x_y_yaw_vel_history_np[:, i]
                        )
                        self.sensor_history_handle[i].set_xdata(
                            range(len(sensor_x_y_yaw_vel_history_np))
                        )
                        self.sensor_history_handle[i].set_ydata(
                            sensor_x_y_yaw_vel_history_np[:, i]
                        )
                        self.actual_history_handle[i].set_xdata(
                            range(len(actual_x_y_yaw_vel_history_np))
                        )
                        self.actual_history_handle[i].set_ydata(
                            actual_x_y_yaw_vel_history_np[:, i]
                        )
                        self.axis[i].relim()
                        self.axis[i].autoscale_view()

                        var_x = np.var(
                            actual_x_y_yaw_vel_history_np[:, 0]
                            - sensor_x_y_yaw_vel_history_np[:, 0]
                        )
                        var_y = np.var(
                            actual_x_y_yaw_vel_history_np[:, 1]
                            - sensor_x_y_yaw_vel_history_np[:, 1]
                        )
                        error_ekf_x = np.mean(
                            np.abs(
                                actual_x_y_yaw_vel_history_np[:, 0]
                                - ekf_x_y_yaw_vel_history_np[:, 0]
                            )
                        )
                        error_ekf_y = np.mean(
                            np.abs(
                                actual_x_y_yaw_vel_history_np[:, 1]
                                - ekf_x_y_yaw_vel_history_np[:, 1]
                            )
                        )
                        self.axis[0].set_title(
                            f"var_x: {var_x:.3f}, var_y: {var_y:.3f}, error_ekf_x: {error_ekf_x:.3f}, error_ekf_y: {error_ekf_y:.3f}",
                            fontsize=11,
                        )

                    plt.pause(1e-3)

            return DoraStatus.CONTINUE
